import matplotlib.pyplot as plt

def visualize_loss(history, title, show=False):
    assert isinstance(history, dict)

    COLORS='brgp'
    plt.figure(figsize=(10, 5))
    for i, (key, value) in enumerate( history.items()):
        color = COLORS[i]
        epochs = range(len(value))
        plt.plot(epochs, value, color, label=key)

    plt.title(title)
    plt.xlabel("Epochs")
    plt.ylabel("Value")
    plt.legend()
    if show:
        plt.show()
    plt.savefig(title+'.png')


class TwinAxisPlotter:
    """
    画一个loss和bleu同框并且趋势都比较明显的图。
    """
    BLEU_SCALE=100
    LOSS_SCALE=1

    def __init__(self, loss_list, bleu_list, label) -> None:
        self.loss_list=[x * self.LOSS_SCALE for x in loss_list]
        self.bleu_list=[x * self.BLEU_SCALE for x in bleu_list]
        self.label = label

    def plot(self, output_file, title):
        plt.figure(dpi=300)
        x_list = range(0, len(self.loss_list))
        plt.figure()
        plt.grid(ls='--')
        ax1 = plt.subplot()
        ax2 = ax1.twinx()

        ln1 = ax1.plot(x_list, self.loss_list, color='C0', label='Loss')
        ax1.set_ylabel('Loss')
        ax1.set_xlabel('Iteration number')

        ln2 = ax2.plot(x_list, self.bleu_list, color='C1', label=self.label)
        ax2.set_ylabel(f'{self.label} (%)')
        lns = ln1 + ln2

        plt.legend(lns, [l.get_label() for l in lns], loc='best')
        plt.title(title)
        plt.savefig(output_file)
        plt.close()
        print(output_file)


if __name__ == "__main__":
    import numpy as np

    x=np.arange(100)

    obj=TwinAxisPlotter(
        loss_list=x,
        bleu_list=-x,
    )
    obj.plot("loss-bleu.png")
